package com.jxpanda.starter.config.security.authentication;

import com.jxpanda.commons.constant.StringConstant;
import com.jxpanda.commons.toolkit.ObjectKit;
import com.jxpanda.commons.toolkit.StringKit;
import com.jxpanda.starter.config.security.Device;
import com.jxpanda.starter.config.security.token.AbstractTokenHelper;
import com.jxpanda.starter.config.security.token.Token;
import io.jsonwebtoken.lang.Collections;
import lombok.AllArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.springframework.http.HttpHeaders;
import org.springframework.http.server.reactive.ServerHttpRequest;
import org.springframework.security.authentication.ReactiveAuthenticationManager;
import org.springframework.security.core.context.SecurityContext;
import org.springframework.security.core.context.SecurityContextImpl;
import org.springframework.security.web.server.context.ServerSecurityContextRepository;
import org.springframework.util.MultiValueMap;
import org.springframework.util.MultiValueMapAdapter;
import org.springframework.web.server.ServerWebExchange;
import reactor.core.publisher.Mono;

import java.lang.reflect.Field;
import java.util.Arrays;
import java.util.HashMap;
import java.util.Objects;

/**
 * @author Panda
 */
@Slf4j
@AllArgsConstructor
public class TokenSecurityContextRepository implements ServerSecurityContextRepository {

    private static final String USER_ID = "userId";
    private static final String STAFF_ID = "staffId";
    private static final String IP = "ip";
    private static final String DEVICE = "device";
    private static final String COORDINATE = "coordinate";
    public static final String TARGET_MAP = "targetMap";

    private final AbstractTokenHelper tokenHelper;

    private final ReactiveAuthenticationManager authenticationManager;

    @Override
    public Mono<Void> save(ServerWebExchange exchange, SecurityContext context) {
        return Mono.empty();
    }

    @Override
    public Mono<SecurityContext> load(ServerWebExchange exchange) {
        ServerHttpRequest request = exchange.getRequest();
        String authorization = getHeader(request.getHeaders(), HttpHeaders.AUTHORIZATION);
        Token token = this.tokenHelper.validate(authorization);
        injectToken(exchange, token);
        return this.authenticationManager
                .authenticate(new TokenAuthentication(token))
                .map(SecurityContextImpl::new);
    }

    /**
     * 从headers中获取header
     *
     * @param headers 请求的headers
     * @param key     要获取的header对应的key
     * @return header中的值
     */
    private static String getHeader(HttpHeaders headers, String key) {
        String first = headers.getFirst(key);
        return ObjectKit.isEmpty(first) ? StringConstant.BLANK : first;
    }


    public static void injectToken(ServerWebExchange exchange, Token token) {
        ServerHttpRequest request = exchange.getRequest();

        // 注入设备标识符、IP等常用属性
        MultiValueMap<String, Object> injectParam = new MultiValueMapAdapter<>(new HashMap<>(8));

        Device device = Arrays.stream(Device.values())
                .filter(it -> it.name().equals(getHeader(request.getHeaders(), DEVICE)))
                .findFirst()
                .orElse(Device.UNKNOWN);
        injectParam.add(DEVICE, device);
        injectParam.add(IP, Objects.requireNonNull(request.getRemoteAddress()).getAddress().getHostAddress());
        // 注入坐标
        String coordinate = getHeader(request.getHeaders(), COORDINATE);
        if (StringKit.isNotBlank(coordinate)) {
            injectParam.add(COORDINATE, coordinate);
        }

        if (token.isEffective()) {
            // 注入用户ID，方便后面接口取值
            injectParam.add(USER_ID, token.getId());
            injectParam.add(STAFF_ID, token.getId());
            // 扩展参数注入，这个参数负责透传后端给出去的值
            if (!Collections.isEmpty(token.getExtra())) {
                token.getExtra().forEach(injectParam::add);
            }
        }

        try {
            // 使用反射的方式替换原来request中的参数
            // 因为原request中的参数是不允许球盖的（调用 add,put 等函数会报错）
            // 所以需要外部创建一个能修改的，然后用反射替换掉原来的
            MultiValueMap<String, String> queryParams = request.getQueryParams();
            // 入参合并
            queryParams.forEach(injectParam::addIfAbsent);

            // 看源码得知，MultiValueMap对象的数据是放在 targetMap 属性中的
            Field targetMap = queryParams.getClass().getDeclaredField(TARGET_MAP);
            // 强行可访问
            targetMap.setAccessible(true);
            // 将原有参数表替换
            targetMap.set(queryParams, injectParam);
        } catch (Exception e) {
            log.error("[REQUEST PARAM INJECT ERROR]", e);
        }
    }

}
